# -*- coding: utf-8 -*-

import numpy as np
from scipy.signal import hilbert
from scipy.fft import fft

def compute_frequency_centers(signal, fs, num_segments):
    spectrum = fft(signal)
    freqs = np.fft.fftfreq(len(spectrum), 1/fs)
    segment_size = len(freqs) // num_segments
    # Compute frequency centers for each segment
    centers = []
    for i in range(num_segments):
        start_idx = i * segment_size
        end_idx = (i + 1) * segment_size
        segment_spectrum = spectrum[start_idx:end_idx]
        segment_freqs = freqs[start_idx:end_idx]
        center_freq = np.sum(np.abs(segment_spectrum) * segment_freqs) / np.sum(np.abs(segment_spectrum))
        centers.append(center_freq)
    
    return centers

def determine_vmd_modes(signal, fs, num_segments, threshold=0.5):
    previous_centers = None
    num_modes = 0
    for k in range(1, num_segments+1):
        analytic_signal = hilbert(signal)
        frequency_centers = compute_frequency_centers(np.abs(analytic_signal), fs, k)     
        # Check if the center frequencies of adjacent modes are close
        if previous_centers is not None:
            center_overlap = np.mean(np.abs(np.array(frequency_centers) - np.array(previous_centers)))
            if center_overlap < threshold:
                # If overlap is less than threshold, break the loop
                break
        
        # Update variables for the next iteration
        previous_centers = frequency_centers
        num_modes = k  
    return num_modes


